# optimization/optimizer.py
from typing import Type, Dict, Any, List, Tuple, Callable, TYPE_CHECKING, Optional  # 导入 TYPE_CHECKING
import pandas as pd
from tqdm import tqdm

# from engine import Operator # <--- 移除这个直接导入
from strategy.strategy import Strategy
from utils.klines_loader import KLines
from utils.parameter_space import ParameterGrid

if TYPE_CHECKING:  # 仅用于类型检查，避免运行时循环导入
    from engine.engine import Operator  # 导入 Operator 以便类型提示


class GridSearchOptimizer:
    def __init__(
            self,
            operator_class: Type['Operator'],  # <--- 修改点: 接收 Operator 类
            execution_klines: KLines,
            other_tf_klines_map: Optional[Dict[str, KLines]],
            strategy_class: Type[Strategy],
            param_grid_dict: Dict[str, List[Any]],
            initial_balance: float = 10000.0,
            open_commission_rate: float = 0.0008,
            close_commission_rate: float = 0.001,
            slippage_percent: float = 0.0005,
            leverage: float = 1.0,
            evaluation_metric: Callable[[Dict[str, Any]], float] = lambda p: p.get("total_net_profit", 0.0)
    ):
        self.operator_class = operator_class  # <--- 存储 Operator 类
        self.execution_klines = execution_klines
        # ... (其余 __init__ 内容不变) ...
        self.other_tf_klines_map = other_tf_klines_map
        self.strategy_class = strategy_class
        self.param_grid = ParameterGrid(param_grid_dict)
        self.evaluation_metric = evaluation_metric

        self.operator_fixed_params = {
            "initial_balance": initial_balance,
            "open_commission_rate": open_commission_rate,
            "close_commission_rate": close_commission_rate,
            "slippage_percent": slippage_percent,
            "leverage": leverage,
        }
        self.results: List[Dict[str, Any]] = []

    def run_optimization(self) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
        """
        执行网格搜索优化。
        :return: 一个元组 (最优参数组合, 最优参数组合对应的性能结果)，如果无结果则为 (None, None)
        """
        self.results = []
        best_score = -float('inf')
        best_params = None
        best_performance = None

        print(f"开始网格搜索优化，总共 {len(self.param_grid)} 组参数...")

        # 使用 tqdm 创建进度条
        for strategy_params_set in tqdm(self.param_grid, desc="优化进度"):
            print(f"\n测试参数: {strategy_params_set}")
            try:
                # 修改点: 使用 self.operator_class 来实例化
                operator_instance = self.operator_class(
                    execution_klines_data=self.execution_klines,
                    other_tf_klines_map=self.other_tf_klines_map,
                    strategy_class=self.strategy_class,
                    strategy_params=strategy_params_set,
                    **self.operator_fixed_params
                )
                performance = operator_instance.run_backtest()

                current_score = self.evaluation_metric(performance)

                self.results.append({
                    "params": strategy_params_set,
                    "performance": performance,  # 存储完整的性能数据
                    "score": current_score
                })

                print(
                    f"参数组: {strategy_params_set}, 得分: {current_score:.4f}, 净利润: {performance.get('total_net_profit', 0):.2f}")

                if current_score > best_score:
                    best_score = current_score
                    best_params = strategy_params_set
                    best_performance = performance

            except Exception as e:
                print(f"参数 {strategy_params_set} 回测时出错: {e}")
                import traceback
                traceback.print_exc()
                self.results.append({
                    "params": strategy_params_set,
                    "performance": None,
                    "score": -float('inf'),  # 标记为失败
                    "error": str(e)
                })

        if best_params:
            print(f"\n优化完成。最优参数: {best_params}, 最高得分: {best_score:.4f}")
        else:
            print("\n优化完成，但未找到有效结果。")

        return best_params, best_performance

    def get_results_dataframe(self) -> pd.DataFrame:
        """将优化结果转换为DataFrame，方便分析。"""
        if not self.results:
            return pd.DataFrame()

        # 将嵌套的 params 和 performance 字典展平，以便更容易放入DataFrame
        flat_results = []
        for res in self.results:
            row = {}
            if res["params"]:
                for p_name, p_val in res["params"].items():
                    row[f"param_{p_name}"] = p_val

            if res["performance"]:
                # 只选择一些关键性能指标放入DataFrame，避免过于庞大
                for metric_name, metric_val in res["performance"].items():
                    if metric_name not in ["equity_curve", "trades_history"]:  # 排除大型列表
                        if isinstance(metric_val, (int, float, str, bool)):  # 只取简单类型
                            row[metric_name] = metric_val

            row["score"] = res.get("score")
            row["error"] = res.get("error")
            flat_results.append(row)

        return pd.DataFrame(flat_results)